#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Word文档提取器 - 核心提取逻辑
个人剂量监测报告数据提取的核心实现
"""

import os
import re
from docx import Document
import json
import csv
from pathlib import Path

from core.data_processor import (
    process_dose_value, get_protection_type, 
    map_occupation_code, check_dose_level
)
from core.excel_lookup import ExcelLookup
from utils.helpers import parse_address, format_date_for_display
from app.config import CSV_FIELDNAMES


class DocumentExtractor:
    """Word文档提取器"""
    
    def __init__(self, excel_file_path=None):
        """初始化提取器
        
        Args:
            excel_file_path: Excel查询文件路径
        """
        self.data = {}
        self.tables = []
        
        # 初始化Excel查询模块
        self.excel_lookup = ExcelLookup(excel_file_path)
    
    def extract_table_data(self, doc):
        """提取表格数据
        
        Args:
            doc: python-docx Document对象
            
        Returns:
            list: 提取的表格数据
        """
        tables = []
        for table in doc.tables:
            table_data = []
            for row in table.rows:
                row_data = []
                for cell in row.cells:
                    row_data.append(cell.text.strip())
                table_data.append(row_data)
            tables.append(table_data)
        return tables
    
    def extract_unit_info(self, table1):
        """从表格1提取单位信息
        
        Args:
            table1: 第一个表格的数据
            
        Returns:
            tuple: (单位名称, 单位地址, 样品数量, 检测仪器)
        """
        unit_name = ""
        unit_address = ""
        sample_count = ""
        detection_device = ""
        
        try:
            # 表格1行1: 单位名称
            if len(table1) > 0 and len(table1[0]) > 0:
                row1_text = table1[0][0]
                # 使用DOTALL标志处理跨行的单位名称，并清理换行符和空格
                name_match = re.search(r'单位名称：(.+?)(?:\s*\||$)', row1_text, re.DOTALL)
                if name_match:
                    # 清理单位名称文本：移除多余的换行符和空格
                    unit_name = re.sub(r'\s+', ' ', name_match.group(1).strip())
            
            # 表格1行2: 单位地址
            if len(table1) > 1 and len(table1[1]) > 0:
                row2_text = table1[1][0]
                # 使用DOTALL标志处理跨行的地址文本，并清理换行符和空格
                address_match = re.search(r'单位地址：(.+?)(?:\s*\||$)', row2_text, re.DOTALL)
                if address_match:
                    # 清理地址文本：移除多余的换行符和空格
                    unit_address = re.sub(r'\s+', ' ', address_match.group(1).strip())
            
            # 表格1行4: 检测仪器
            if len(table1) > 3 and len(table1[3]) > 0:
                row4_text = table1[3][0]
                device_match = re.search(r'检测仪器：(.+?)(?:\s*\||$)', row4_text)
                if device_match:
                    detection_device = device_match.group(1).strip()
            
            # 表格1行5: 样品数量
            if len(table1) > 4:
                row5_text = ""
                if len(table1[4]) > 1:
                    row5_text = table1[4][1]
                elif len(table1[4]) > 0:
                    row5_text = table1[4][0]
                
                sample_match = re.search(r'样品数量：(.+?)(?:\s*\||$)', row5_text)
                if sample_match:
                    sample_count = sample_match.group(1).strip()
        
        except Exception as e:
            print(f"提取单位信息出错: {e}")
        
        return unit_name, unit_address, sample_count, detection_device
    
    def extract_task_number(self, table2):
        """从表格2提取检测任务编号
        
        Args:
            table2: 第二个表格的数据
            
        Returns:
            str: 检测任务编号
        """
        task_number = ""
        try:
            # 遍历表格2的所有单元格查找检测任务编号
            for row in table2:
                for cell in row:
                    if "检测任务编号" in cell:
                        task_match = re.search(r'检测任务编号：\s*([A-Z0-9]+)', cell)
                        if task_match:
                            task_number = task_match.group(1).strip()
                            break
                if task_number:
                    break
        except Exception as e:
            print(f"提取检测任务编号出错: {e}")
        
        return task_number
    
    def extract_header_info(self, doc):
        """从个人剂量监测报告文档页眉提取收样日期、报告编号、检测日期
        
        Args:
            doc: python-docx Document对象
            
        Returns:
            tuple: (收样日期, 报告编号, 检测日期)
        """
        receive_date = ""
        report_number = ""
        test_date = ""
        
        try:
            # 遍历所有section的页眉
            for section in doc.sections:
                header = section.header
                if header:
                    # 检查页眉中的段落
                    for para in header.paragraphs:
                        text = para.text.strip()
                        if text:
                            # 在段落文本中也尝试提取信息
                            if '收样日期' in text and not receive_date:
                                receive_match = re.search(r'收样日期[：:]\s*(\d{4}\s*年\s*\d{1,2}\s*月\s*\d{1,2}\s*日)', text)
                                if receive_match:
                                    receive_date = receive_match.group(1).strip()
                            
                            if '报告编号' in text and not report_number:
                                report_match = re.search(r'报告编号[：:]\s*([A-Z]{2,}[-_]?[A-Z0-9]+)', text)
                                if report_match:
                                    report_number = report_match.group(1).strip()
                            
                            if '检测日期' in text and not test_date:
                                test_match = re.search(r'检测日期[：:]\s*(\d{4}\s*年\s*\d{1,2}\s*月\s*\d{1,2}\s*日)', text)
                                if test_match:
                                    test_date = test_match.group(1).strip()
                    
                    # 检查页眉中的表格
                    for table in header.tables:
                        for row in table.rows:
                            for cell in row.cells:
                                text = cell.text.strip()
                                if text:
                                    # 提取收样日期
                                    if '收样日期' in text and not receive_date:
                                        receive_match = re.search(r'收样日期[：:]\s*(\d{4}\s*年\s*\d{1,2}\s*月\s*\d{1,2}\s*日)', text)
                                        if receive_match:
                                            receive_date = receive_match.group(1).strip()
                                    
                                    # 提取报告编号
                                    if '报告编号' in text and not report_number:
                                        report_match = re.search(r'报告编号[：:]\s*([A-Z]{2,}[-_]?[A-Z0-9]+)', text)
                                        if report_match:
                                            report_number = report_match.group(1).strip()
                                    
                                    # 提取检测日期
                                    if '检测日期' in text and not test_date:
                                        test_match = re.search(r'检测日期[：:]\s*(\d{4}\s*年\s*\d{1,2}\s*月\s*\d{1,2}\s*日)', text)
                                        if test_match:
                                            test_date = test_match.group(1).strip()
                    
                    # 如果找到了所有信息就退出
                    if receive_date and report_number and test_date:
                        break
        
        except Exception as e:
            print(f"提取页眉信息出错: {e}")
        
        return receive_date, report_number, test_date

    def extract_personnel_info(self, table4):
        """从表格4提取所有人员信息
        
        Args:
            table4: 第四个表格的数据
            
        Returns:
            list: 包含所有人员信息的列表，每个元素是(部门, 样品编号, 姓名, 职业类别, 佩戴周期, 剂量当量)
        """
        all_personnel = []
        current_department = ""
        
        try:
            # 跳过表头（第一行），从第二行开始处理
            for row_index in range(1, len(table4)):
                row = table4[row_index]
                
                if not row or len(row) == 0:
                    continue
                
                # 检查是否为部门标题行（所有列都填充相同的部门名称）
                first_cell = row[0].strip() if row[0] else ""
                is_department_row = False
                
                # 判断部门行的条件：
                # 1. 第一列有内容
                # 2. 所有列的内容都相同（都是部门名称）
                # 3. 内容不像序号（不是纯数字）且不是JL开头的样品编号
                # 4. 不是表头（不包含"序号"、"编号"、"姓名"等）
                if (first_cell and 
                    not first_cell.isdigit() and 
                    not (first_cell.startswith('JL') or 
                         (first_cell.startswith('J') and len(first_cell) > 1 and 
                          any(c.isdigit() for c in first_cell[1:4]))) and
                    first_cell not in ['序号', '编号', '姓名', '职业类别', '剂量计佩戴周期', '剂量当量']):
                    
                    # 检查所有列是否都是相同的内容（部门名称）
                    all_cells_same = True
                    non_empty_count = 0
                    for i in range(0, min(len(row), 6)):
                        if row[i] and row[i].strip():
                            non_empty_count += 1
                            if row[i].strip() != first_cell:
                                all_cells_same = False
                                break
                    
                    # 至少要有2列以上相同内容，且所有非空列都相同
                    if all_cells_same and non_empty_count >= 2:
                        is_department_row = True
                
                if is_department_row:
                    # 这是部门标题行，更新当前部门
                    if "|" in first_cell:
                        current_department = first_cell.split("|")[0].strip()
                    else:
                        current_department = first_cell.strip()
                    print(f"发现部门: {current_department}")
                    continue
                
                # 检查是否为人员数据行（需要有序号、样品编号、姓名等）
                if len(row) >= 6:
                    # 尝试按列提取数据
                    sequence = row[0].strip() if row[0] else ""
                    sample_id = row[1].strip() if row[1] else ""
                    name = row[2].strip() if row[2] else ""
                    protection_type = row[3].strip() if row[3] else ""
                    wearing_period = row[4].strip() if row[4] else ""
                    dose_value = row[5].strip() if row[5] else ""
                    
                    # 验证这是一个有效的人员记录
                    # 1. 至少要有姓名和样品编号
                    # 2. 样品编号应该是JL开头的格式，不是部门名称
                    # 3. 序号应该是数字
                    is_valid_person = (
                        name and sample_id and
                        (sample_id.startswith('JL') or 
                         (sample_id.startswith('J') and len(sample_id) > 1 and 
                          any(c.isdigit() for c in sample_id[1:4]))) and
                        sequence.isdigit()
                    )
                    
                    if is_valid_person:
                        all_personnel.append((
                            current_department,
                            sample_id,
                            name,
                            protection_type,
                            wearing_period,
                            dose_value
                        ))
                        print(f"提取人员: {current_department} - {name} ({sample_id})")
                
                elif len(row) >= 1:
                    # 尝试处理用|分割的格式
                    row_text = row[0]
                    if "|" in row_text:
                        parts = row_text.split("|")
                        if len(parts) >= 6:
                            sequence = parts[0].strip()
                            sample_id = parts[1].strip()
                            name = parts[2].strip()
                            protection_type = parts[3].strip()
                            wearing_period = parts[4].strip()
                            dose_value = parts[5].strip()
                            
                            # 验证这是一个有效的人员记录
                            is_valid_person = (
                                name and sample_id and
                                (sample_id.startswith('JL') or 
                         (sample_id.startswith('J') and len(sample_id) > 1 and 
                          any(c.isdigit() for c in sample_id[1:4]))) and
                                sequence.isdigit()
                            )
                            
                            if is_valid_person:
                                all_personnel.append((
                                    current_department,
                                    sample_id,
                                    name,
                                    protection_type,
                                    wearing_period,
                                    dose_value
                                ))
                                print(f"提取人员: {current_department} - {name} ({sample_id})")
        
        except Exception as e:
            print(f"提取人员信息出错: {e}")
        
        print(f"总共提取到 {len(all_personnel)} 个人员")
        return all_personnel
    
    def extract_personnel_info_with_context(self, table_data, inherited_department=""):
        """从表格数据提取人员信息，支持跨表格部门延续
        
        Args:
            table_data: 表格数据
            inherited_department: 从上一个表格继承的部门信息
            
        Returns:
            list: 包含所有人员信息的列表，每个元素是(部门, 样品编号, 姓名, 职业类别, 佩戴周期, 剂量当量)
        """
        all_personnel = []
        current_department = inherited_department  # 从继承的部门开始
        
        try:
            # 跳过表头（第一行），从第二行开始处理
            for row_index in range(1, len(table_data)):
                row = table_data[row_index]
                
                if not row or len(row) == 0:
                    continue
                
                # 检查是否为部门标题行（所有列都填充相同的部门名称）
                first_cell = row[0].strip() if row[0] else ""
                is_department_row = False
                
                # 判断部门行的条件：
                # 1. 第一列有内容
                # 2. 所有列的内容都相同（都是部门名称）
                # 3. 内容不像序号（不是纯数字）且不是JL开头的样品编号
                # 4. 不是表头（不包含"序号"、"编号"、"姓名"等）
                if (first_cell and 
                    not first_cell.isdigit() and 
                    not (first_cell.startswith('JL') or 
                         (first_cell.startswith('J') and len(first_cell) > 1 and 
                          any(c.isdigit() for c in first_cell[1:4]))) and
                    first_cell not in ['序号', '编号', '姓名', '职业类别', '剂量计佩戴周期', '剂量当量']):
                    
                    # 检查所有列是否都是相同的内容（部门名称）
                    all_cells_same = True
                    non_empty_count = 0
                    for i in range(0, min(len(row), 6)):
                        if row[i] and row[i].strip():
                            non_empty_count += 1
                            if row[i].strip() != first_cell:
                                all_cells_same = False
                                break
                    
                    # 至少要有2列以上相同内容，且所有非空列都相同
                    if all_cells_same and non_empty_count >= 2:
                        is_department_row = True
                
                if is_department_row:
                    # 这是部门标题行，更新当前部门
                    if "|" in first_cell:
                        current_department = first_cell.split("|")[0].strip()
                    else:
                        current_department = first_cell.strip()
                    print(f"发现部门: {current_department}")
                    continue
                
                # 检查是否为人员数据行（需要有序号、样品编号、姓名等）
                if len(row) >= 6:
                    # 尝试按列提取数据
                    sequence = row[0].strip() if row[0] else ""
                    sample_id = row[1].strip() if row[1] else ""
                    name = row[2].strip() if row[2] else ""
                    protection_type = row[3].strip() if row[3] else ""
                    wearing_period = row[4].strip() if row[4] else ""
                    dose_value = row[5].strip() if row[5] else ""
                    
                    # 验证这是一个有效的人员记录
                    # 1. 至少要有姓名和样品编号
                    # 2. 样品编号应该是JL开头的格式，不是部门名称
                    # 3. 序号应该是数字
                    is_valid_person = (
                        name and sample_id and
                        (sample_id.startswith('JL') or 
                         (sample_id.startswith('J') and len(sample_id) > 1 and 
                          any(c.isdigit() for c in sample_id[1:4]))) and
                        sequence.isdigit()
                    )
                    
                    if is_valid_person:
                        # 如果当前没有部门信息，但有继承的部门信息，使用继承的部门
                        final_department = current_department if current_department else inherited_department
                        
                        all_personnel.append((
                            final_department,
                            sample_id,
                            name,
                            protection_type,
                            wearing_period,
                            dose_value
                        ))
                        print(f"提取人员: {final_department} - {name} ({sample_id})")
                
                elif len(row) >= 1:
                    # 尝试处理用|分割的格式
                    row_text = row[0]
                    if "|" in row_text:
                        parts = row_text.split("|")
                        if len(parts) >= 6:
                            sequence = parts[0].strip()
                            sample_id = parts[1].strip()
                            name = parts[2].strip()
                            protection_type = parts[3].strip()
                            wearing_period = parts[4].strip()
                            dose_value = parts[5].strip()
                            
                            # 验证这是一个有效的人员记录
                            is_valid_person = (
                                name and sample_id and
                                (sample_id.startswith('JL') or 
                         (sample_id.startswith('J') and len(sample_id) > 1 and 
                          any(c.isdigit() for c in sample_id[1:4]))) and
                                sequence.isdigit()
                            )
                            
                            if is_valid_person:
                                # 如果当前没有部门信息，但有继承的部门信息，使用继承的部门
                                final_department = current_department if current_department else inherited_department
                                
                                all_personnel.append((
                                    final_department,
                                    sample_id,
                                    name,
                                    protection_type,
                                    wearing_period,
                                    dose_value
                                ))
                                print(f"提取人员: {final_department} - {name} ({sample_id})")
        
        except Exception as e:
            print(f"提取人员信息出错: {e}")
        
        print(f"表格提取到 {len(all_personnel)} 个人员")
        return all_personnel
    
    def extract_from_file(self, file_path):
        """从Word文件提取所有需要的数据
        
        Args:
            file_path: Word文档路径
            
        Returns:
            list: 包含所有人员数据字典的列表
        """
        try:
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件不存在: {file_path}")
            
            if not file_path.endswith('.docx'):
                raise ValueError("仅支持.docx格式的文件")
            
            doc = Document(file_path)
            self.tables = self.extract_table_data(doc)
            
            # 检查是否有足够的表格
            if len(self.tables) < 4:
                print(f"警告: 文档中只有 {len(self.tables)} 个表格，可能无法提取所有信息")
            
            # 提取各种信息
            unit_name, unit_address, sample_count, detection_device = self.extract_unit_info(self.tables[0])
            
            # 解析地址
            province, city, district, postal_code = parse_address(unit_address)
            
            # 提取头部信息（收样日期、报告编号、检测日期）
            receive_date, report_number, test_date = self.extract_header_info(doc)
            
            # 提取检测任务编号
            task_number = ""
            if len(self.tables) > 1:
                task_number = self.extract_task_number(self.tables[1])
            
            # 从所有表格中提取人员信息，支持跨表格部门延续
            all_personnel = []
            last_department = ""  # 用于跨表格部门延续
            
            for table_index, table in enumerate(self.tables):
                if len(table) > 1:  # 至少有表头和一行数据
                    # 检查这个表格是否包含人员信息
                    has_personnel = False
                    for row in table[:5]:  # 检查前5行
                        if len(row) >= 6:
                            # 检查是否有样品编号（JL开头或J开头+数字）
                            for cell in row:
                                if cell and isinstance(cell, str):
                                    cell_text = cell.strip()
                                    # 更灵活的样品编号识别：JL开头 或 J+数字开头
                                    if (cell_text.startswith('JL') or 
                                        (cell_text.startswith('J') and len(cell_text) > 1 and 
                                         any(c.isdigit() for c in cell_text[1:4]))):  # J后面有数字
                                        has_personnel = True
                                        break
                        if has_personnel:
                            break
                    
                    if has_personnel:
                        print(f"在表格{table_index + 1}中发现人员信息，开始提取...")
                        table_personnel = self.extract_personnel_info_with_context(table, last_department)
                        all_personnel.extend(table_personnel)
                        
                        # 更新最后一个部门信息，用于下一个表格
                        if table_personnel:
                            last_department = table_personnel[-1][0]  # 获取最后一个人员的部门信息
            
            # 如果没有提取到人员信息，返回空列表
            if not all_personnel:
                print(f"警告: 文档 {file_path} 中没有提取到人员信息")
                return []
            
            # 为每个人员创建完整的数据记录
            all_records = []
            for department, sample_id, name, original_occupation_code, wearing_period, dose_value in all_personnel:
                # 处理剂量当量值
                processed_dose = process_dose_value(dose_value)
                
                # 从Excel中查找人员信息
                gender, id_card = self.excel_lookup.lookup_person_info(sample_id)
                
                # 映射职业类别
                occupation_description = map_occupation_code(original_occupation_code)
                
                # 确定防护类型
                protection_type = get_protection_type(original_occupation_code)
                
                # 检查是否超调查水平
                dose_level_exceeded = check_dose_level(processed_dose)
                
                # 组装数据
                record = {
                    'file_path': file_path,
                    '受检单位名称': unit_name,
                    '单位地址': unit_address,
                    '省': province,
                    '市': city,
                    '区（镇）': district,
                    '邮政编码': postal_code,
                    '收样日期': format_date_for_display(receive_date),
                    '报告编号': report_number,
                    '检测日期': format_date_for_display(test_date),
                    '样品数量': sample_count,
                    '检测仪器': detection_device,
                    '检测任务编号': task_number,
                    '部门': department,
                    '样品编号': sample_id,
                    '姓名': name,
                    '身份证号': id_card,
                    '性别': gender,
                    '射线及防护类型（是否穿铅衣）': protection_type,
                    '职业类别': occupation_description,
                    '佩戴周期': wearing_period,
                    '剂量当量': processed_dose,
                    '是否超调查水平（1.25mSv）': dose_level_exceeded
                }
                all_records.append(record)
            
            print(f"从文件 {file_path} 中提取到 {len(all_records)} 个人员记录")
            return all_records
            
        except Exception as e:
            raise Exception(f"提取文档内容失败: {str(e)}")
    
    def save_to_json(self, output_path):
        """保存为JSON格式
        
        Args:
            output_path: 输出文件路径
        """
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(self.data, f, ensure_ascii=False, indent=2)
    
    def save_to_csv(self, output_path):
        """保存为CSV格式
        
        Args:
            output_path: 输出文件路径
        """
        if not self.data:
            raise ValueError("没有数据可保存")
        
        # 创建只包含指定字段的字典
        csv_data = {field: self.data.get(field, '') for field in CSV_FIELDNAMES}
        
        with open(output_path, 'w', newline='', encoding='utf-8-sig') as f:
            writer = csv.DictWriter(f, fieldnames=CSV_FIELDNAMES)
            writer.writeheader()
            writer.writerow(csv_data)
    
    def get_debug_info(self):
        """获取调试信息
        
        Returns:
            dict: 调试信息字典
        """
        debug_info = {
            'total_tables': len(self.tables),
            'table_structures': []
        }
        
        for i, table in enumerate(self.tables):
            table_info = {
                'table_index': i,
                'rows': len(table),
                'structure': []
            }
            for j, row in enumerate(table):
                table_info['structure'].append({
                    'row_index': j,
                    'columns': len(row),
                    'content': row
                })
            debug_info['table_structures'].append(table_info)
        
        return debug_info
